# 安装必要的库
# %pip install langchain langchain-community langchain-openai  langgraph faiss-cpu sqlite3

import os
import getpass
from langchain_community.utilities import SQLDatabase
# from langchain_openai import ChatOpenAI
from langchain.chains import create_sql_query_chain
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_community.tools.sql_database.tool import QuerySQLDataBaseTool
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_core.messages import SystemMessage
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent
from langchain_core.messages import HumanMessage
from langchain_community.chat_models import ChatZhipuAI
# 去智普申请
print(os.environ["ZHIPUAI_API_KEY"])
# 初始化数据库并插入数据
db = SQLDatabase.from_uri("sqlite:///Chinook.db")

# 检查表是否已经存在
table_check_sql = '''
SELECT name FROM sqlite_master WHERE type='table' AND name='Employee';
'''
table_exists = db.run(table_check_sql)

# 如果表不存在，则创建表并插入数据
if not table_exists:
    # 执行建表语句
    create_table_sql = '''
    CREATE TABLE Employee (
        EmployeeId INTEGER NOT NULL,
        LastName NVARCHAR(20) NOT NULL,
        FirstName NVARCHAR(20) NOT NULL,
        Title NVARCHAR(30),
        ReportsTo INTEGER,
        BirthDate DATETIME,
        HireDate DATETIME,
        Address NVARCHAR(70),
        City NVARCHAR(40),
        State NVARCHAR(40),
        Country NVARCHAR(40),
        PostalCode NVARCHAR(10),
        Phone NVARCHAR(24),
        Fax NVARCHAR(24),
        Email NVARCHAR(60),
        CONSTRAINT PK_Employee PRIMARY KEY (EmployeeId),
        FOREIGN KEY (ReportsTo) REFERENCES Employee (EmployeeId) 
            ON DELETE NO ACTION ON UPDATE NO ACTION
    );
    '''

    db.run(create_table_sql)

    # 插入初始化数据
    insert_data_sql = '''
    INSERT INTO [Employee] ([EmployeeId], [LastName], [FirstName], [Title], [ReportsTo], [BirthDate], [HireDate], [Address], [City], [State], [Country], [PostalCode], [Phone], [Fax], [Email]) VALUES
        (1, 'Adams', 'Andrew', 'General Manager', NULL, '1962-02-18', '2002-08-14', '11120 Jasper Ave NW', 'Edmonton', 'AB', 'Canada', 'T5K 2N1', '+1 (780) 428-9482', '+1 (780) 428-3457', 'andrew@chinookcorp.com'),
        (2, 'Edwards', 'Nancy', 'Sales Manager', 1, '1958-12-08', '2002-05-01', '825 8 Ave SW', 'Calgary', 'AB', 'Canada', 'T2P 2T3', '+1 (403) 262-3443', '+1 (403) 262-3322', 'nancy@chinookcorp.com'),
        (3, 'Peacock', 'Jane', 'Sales Support Agent', 2, '1973-08-29', '2002-04-01', '1111 6 Ave SW', 'Calgary', 'AB', 'Canada', 'T2P 5M5', '+1 (403) 262-3443', '+1 (403) 262-6712', 'jane@chinookcorp.com'),
        (4, 'Park', 'Margaret', 'Sales Support Agent', 2, '1947-09-19', '2003-05-03', '683 10 Street SW', 'Calgary', 'AB', 'Canada', 'T2P 5G3', '+1 (403) 263-4423', '+1 (403) 263-4289', 'margaret@chinookcorp.com'),
        (5, 'Johnson', 'Steve', 'Sales Support Agent', 2, '1965-03-03', '2003-10-17', '7727B 41 Ave', 'Calgary', 'AB', 'Canada', 'T3B 1Y7', '1 (780) 836-9987', '1 (780) 836-9543', 'steve@chinookcorp.com'),
        (6, 'Mitchell', 'Michael', 'IT Manager', 1, '1973-07-01', '2003-10-17', '5827 Bowness Road NW', 'Calgary', 'AB', 'Canada', 'T3B 0C5', '+1 (403) 246-9887', '+1 (403) 246-9899', 'michael@chinookcorp.com'),
        (7, 'King', 'Robert', 'IT Staff', 6, '1970-05-29', '2004-01-02', '590 Columbia Boulevard West', 'Lethbridge', 'AB', 'Canada', 'T1K 5N8', '+1 (403) 456-9986', '+1 (403) 456-8485', 'robert@chinookcorp.com'),
        (8, 'Callahan', 'Laura', 'IT Staff', 6, '1968-01-09', '2004-03-04', '923 7 ST NW', 'Lethbridge', 'AB', 'Canada', 'T1H 1Y8', '+1 (403) 467-3351', '+1 (403) 467-8772', 'laura@chinookcorp.com');
    '''

    db.run(insert_data_sql)

# 初始化OpenAI LLM模型
# llm = ChatZhipuAI(model="GLM-4-Flash")

# 初始化语言模型 (可以使用 OpenAI 或其他 LLM)
llm = ChatOpenAI(model='deepseek-ai/DeepSeek-R1').bind(
    logprobs=True)

# ------------------------
# 使用链（Chains）实现问答系统
# ------------------------

# 定义用于回答问题的模板
answer_prompt = PromptTemplate.from_template(
    """Given the following user question, corresponding SQL query, and SQL result, answer the user question.

Question: {question}
SQL Query: {query}
SQL Result: {result}
Answer: """
)

# 将问题转换为SQL查询
write_query = create_sql_query_chain(llm, db)
execute_query = QuerySQLDataBaseTool(db=db)


# 清理查询语句结果
def clean_query(query_result):
    return query_result.get("query", "").replace("SQLQuery:", "").strip()


# 使用 assign + itemgetter 进行数据流转
chain = (
        RunnablePassthrough.assign(query=write_query).assign(
            result=lambda x: execute_query.invoke({"query": clean_query(x)})
        )
        | answer_prompt
        | llm
        | StrOutputParser()
)

# 执行链并回答问题
print(chain.invoke({"question": "How many employees are there"}))

# ------------------------
# 使用代理（Agents）实现问答系统
# ------------------------

# 初始化SQL代理工具包
toolkit = SQLDatabaseToolkit(db=db, llm=llm)
tools = toolkit.get_tools()

# 定义代理系统提示信息
SQL_PREFIX = """You are an agent designed to interact with a SQL database.
Given an input question, create a syntactically correct SQLite query to run, then look at the results of the query and return the answer."""
system_message = SystemMessage(content=SQL_PREFIX)

# 创建代理
agent_executor = create_react_agent(llm, tools, state_modifier=system_message, debug=True)

# 使用代理执行查询并回答问题
for s in agent_executor.stream(
        {"messages": [HumanMessage(content="Who is the oldest employee?")]}
):
    print(s)




